package org.example.utils;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.http.HttpUtil;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import org.example.vo.EmbeddingVO;

import java.math.BigDecimal;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 嵌入模型相关工具类
 */
public class EmbeddingUtil {
    private static final String model = "bge-small-zh-v1.5";
    private static final String url = "http://192.168.1.247:9983/v1/embeddings";

    /**
     * 使用嵌入模型向量化文本
     * @param text 文本
     * @return 向量
     */
    public static double[] getEmbedding(String text) {
        Map<String, Object> requestMap = new HashMap<>();
        requestMap.put("model", model);
        requestMap.put("input", text);
        requestMap.put("encoding_format", "float");

        String result = HttpUtil.post(url, JSONObject.toJSONString(requestMap));
        EmbeddingVO embeddingVO = JSON.parseObject(result, EmbeddingVO.class);
        List<EmbeddingVO.DataInfo> data = embeddingVO.getData();
        if (CollectionUtil.isNotEmpty(data)) {
            return data.get(0).getEmbedding().stream().mapToDouble(BigDecimal::doubleValue).toArray();
        }
        return new double[]{};
    }
}
